# MRI dispersion relation
from numpy import *

Shear = -1.
Omega = 2./3.
Lz    = 1.

def disp(Bz,Nz,Re,Rm):
    kapsq = 2*Omega*(2*Omega+Shear)
    Bsq = Bz**2
    ks = (2*pi/Lz)*arange(int(round(Nz/3.))+1)
    grate = zeros(len(ks))
    if Re==Rm :
        for i in range(len(ks)):
            ksq = ks[i]**2
            c1 = ksq*Bsq +kapsq/2.
            c0 = ksq*Bsq*(ksq*Bsq+2*Omega*Shear)
            if c0 >= 0. :
                grate[i] = -ksq/Rm
            else:
                grate[i] = sqrt(-c0/(c1+sqrt(c1**2-c0))) - ksq/Rm
    else :
        print 'Unable to work with Re != Rm'

    return ks,grate
